"""
ChromaDB 向量数据库管理器
负责文档向量的存储、检索和管理
"""
# 在导入任何模块之前，先禁用ChromaDB的遥测功能
import os
import logging

# 设置环境变量禁用遥测
os.environ["CHROMA_TELEMETRY_ENABLED"] = "false"
os.environ["CHROMA_EXPERIMENTAL_ENABLE_TELEMETRY"] = "false"
os.environ["CHROMA_LOGGING_LEVEL"] = "ERROR"

# 配置Python日志系统，抑制特定错误消息
logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger()

# 添加自定义过滤器来抑制特定的遥测错误消息
class TelemetryErrorFilter(logging.Filter):
    def filter(self, record):
        # 检查日志消息是否包含遥测错误关键字
        if hasattr(record, 'msg'):
            msg = str(record.msg)
            # 抑制包含特定错误模式的日志
            if "Failed to send telemetry event" in msg and "capture() takes 1 positional argument but 3 were given" in msg:
                return False
        return True

# 为所有处理器添加过滤器
for handler in logger.handlers:
    handler.addFilter(TelemetryErrorFilter())

# 同时为ChromaDB的日志记录器添加过滤器
try:
    chroma_logger = logging.getLogger('chromadb')
    for handler in chroma_logger.handlers:
        handler.addFilter(TelemetryErrorFilter())
except Exception:
    pass

# 现在导入其他模块
import uuid
import numpy as np
import shutil
from typing import List, Dict, Any, Optional, Tuple

from langchain.schema import Document
from langchain_community.embeddings import OpenAIEmbeddings, OllamaEmbeddings
from langchain.embeddings.base import Embeddings
from langchain_community.vectorstores import Chroma

# 简单的嵌入模型实现，用于测试和回退
class SimpleEmbeddings(Embeddings):
    """
    改进版简单嵌入模型实现，用于测试和作为最后回退选项
    使用更有意义的方法生成嵌入向量，确保相似文本有相似的嵌入
    """
    
    def __init__(self, dimension: int = 128):
        self.dimension = dimension
        # 常用词的权重映射，用于生成更有意义的嵌入
        self.word_weights = {
            "rag": 0.9, "检索": 0.8, "增强": 0.7, "生成": 0.7,
            "向量": 0.8, "数据库": 0.7, "嵌入": 0.8, "模型": 0.6,
            "llm": 0.8, "知识库": 0.9, "文档": 0.6, "检索": 0.8,
            "相似": 0.5, "搜索": 0.5, "文本": 0.4, "语义": 0.6
        }
    
    def _calculate_text_embedding(self, text: str) -> List[float]:
        """计算单个文本的嵌入向量"""
        # 初始化嵌入向量
        embedding = [0.0] * self.dimension
        
        # 转换文本为小写并分词
        words = text.lower().split()
        
        # 为每个位置分配基于词语的权重
        for i in range(min(len(words), self.dimension)):
            word = words[i]
            # 检查单词是否在权重映射中
            if word in self.word_weights:
                embedding[i] = self.word_weights[word]
            else:
                # 对于不在映射中的词，使用一个基于字符的简单哈希值
                char_sum = sum(ord(c) for c in word)
                embedding[i] = (char_sum % 50) / 100.0 + 0.2  # 范围在0.2-0.7之间
        
        # 确保向量有足够的变化
        for i in range(self.dimension):
            # 添加一些随机性，但保持确定性（基于文本内容）
            seed = hash(text + str(i))
            embedding[i] = (embedding[i] + (seed % 20) / 100.0) % 1.0
        
        return embedding
    
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """为文本列表创建嵌入向量"""
        return [self._calculate_text_embedding(text) for text in texts]
    
    def embed_query(self, text: str) -> List[float]:
        """为查询创建嵌入向量"""
        return self._calculate_text_embedding(text)

class ChromaDBManager:
    """
    ChromaDB 向量数据库管理器
    提供文档向量的存储、检索和管理功能
    """
    
    def __init__(self, collection_name: str = "documents", persist_directory: str = "./vector_db", force_reset: bool = False):
        """
        初始化ChromaDB管理器
        
        Args:
            collection_name: 集合名称
            persist_directory: 持久化目录
            force_reset: 是否强制重置数据库（清理旧数据）
        """
        self.collection_name = collection_name
        self.persist_directory = persist_directory
        self.vectorstore = None
        
        # 如果需要，强制重置数据库（清理旧数据）
        if force_reset and os.path.exists(persist_directory):
            try:
                shutil.rmtree(persist_directory)
                print(f"已清理旧数据库: {persist_directory}")
            except Exception as e:
                print(f"清理旧数据库时出错: {str(e)}")
        
        # 确保持久化目录存在
        os.makedirs(persist_directory, exist_ok=True)
        
        # 获取或创建嵌入模型
        self.embedding_model = self._get_default_embedding_model()
        
        # 初始化向量存储
        self._init_vectorstore()
    
    def _get_default_embedding_model(self) -> Embeddings:
        """
        获取默认的嵌入模型
        按优先级尝试使用OpenAI模型、Ollama模型，如果都失败则使用简单的回退模型
        
        Returns:
            嵌入模型实例
        """
        # 尝试使用OpenAI嵌入模型（如果环境变量设置了API密钥）
        if os.environ.get("OPENAI_API_KEY"):
            try:
                model = OpenAIEmbeddings()
                # 测试模型是否可用
                result = model.embed_query("test")
                print(f"使用OpenAI嵌入模型，嵌入向量维度: {len(result)}")
                return model
            except Exception as e:
                print(f"OpenAI嵌入模型不可用: {str(e)}")
        
        # 尝试使用Ollama嵌入模型
        try:
            print("尝试初始化Ollama嵌入模型...")
            model = OllamaEmbeddings(model="all-minilm:22m")
            # 测试模型是否可用
            print("测试Ollama嵌入模型...")
            result = model.embed_query("test")
            print(f"使用Ollama嵌入模型，嵌入向量维度: {len(result)}")
            
            # 验证两个相似文本的嵌入结果是否相似
            result1 = model.embed_query("向量数据库")
            result2 = model.embed_query("向量存储")
            
            # 计算简单的余弦相似度近似值
            similarity = sum(a*b for a,b in zip(result1[:10], result2[:10]))
            print(f"相似文本嵌入相似度测试: {similarity:.4f}")
            
            return model
        except Exception as e:
            print(f"Ollama嵌入模型不可用或出错: {type(e).__name__}: {str(e)}")
            import traceback
            traceback.print_exc()
        
        # 使用简单的回退模型
        print("警告: 无法使用真实的嵌入模型，使用简单回退模型（仅用于测试）")
        fallback_model = SimpleEmbeddings()
        
        # 测试回退模型
        try:
            result = fallback_model.embed_query("test")
            print(f"使用SimpleEmbeddings模型，嵌入向量维度: {len(result)}")
        except Exception as e:
            print(f"回退模型测试出错: {str(e)}")
            
        return fallback_model
    
    def _init_vectorstore(self):
        """
        初始化向量存储
        如果已存在则加载，否则创建新的
        """
        try:
            # 尝试加载现有向量存储
            self.vectorstore = Chroma(
                persist_directory=self.persist_directory,
                embedding_function=self.embedding_model,
                collection_name=self.collection_name
            )
        except Exception:
            # 创建新的向量存储
            self.vectorstore = Chroma(
                persist_directory=self.persist_directory,
                embedding_function=self.embedding_model,
                collection_name=self.collection_name
            )
    
    def add_documents(self, documents: List[Document]) -> List[str]:
        """
        添加文档到向量数据库
        
        Args:
            documents: 文档列表
            
        Returns:
            添加的文档ID列表
        """
        if not documents:
            return []
        
        try:
            # 添加文档并获取ID
            document_ids = self.vectorstore.add_documents(documents)
            

            
            return document_ids
            
        except Exception as e:
            print(f"添加文档到向量数据库时出错: {str(e)}")
            return []
    
    def add_text(self, text: str, metadata: Optional[Dict[str, Any]] = None) -> Optional[str]:
        """
        添加单个文本到向量数据库
        
        Args:
            text: 文本内容
            metadata: 元数据
            
        Returns:
            添加的文档ID，如果失败则返回None
        """
        if not text:
            return None
        
        metadata = metadata or {}
        document = Document(page_content=text, metadata=metadata)
        
        ids = self.add_documents([document])
        return ids[0] if ids else None
    
    def search(self, query: str, k: int = 5, filter: Optional[Dict[str, Any]] = None) -> List[Tuple[Document, float]]:
        """
        搜索与查询最相似的文档
        
        Args:
            query: 搜索查询
            k: 返回的文档数量
            filter: 过滤条件
            
        Returns:
            文档和相似度分数的列表
        """
        try:
            # 获取实际文档数量，确保k值不超过实际文档数量
            actual_count = self.count_documents()
            safe_k = min(k, actual_count)
            
            # 执行相似度搜索
            results = self.vectorstore.similarity_search_with_score(
                query=query,
                k=safe_k,
                filter=filter
            )
            
            return results
            
        except Exception as e:
            print(f"搜索文档时出错: {str(e)}")
            return []
    
    def get_document(self, document_id: str) -> Optional[Document]:
        """
        通过ID获取文档
        
        Args:
            document_id: 文档ID
            
        Returns:
            文档对象，如果不存在则返回None
        """
        try:
            # 获取文档
            results = self.vectorstore.get(ids=[document_id])
            
            if results and results['documents'] and results['metadata']:
                content = results['documents'][0]
                metadata = results['metadata'][0]
                return Document(page_content=content, metadata=metadata)
            
            return None
            
        except Exception as e:
            print(f"获取文档时出错: {str(e)}")
            return None
    
    def delete_document(self, document_id: str) -> bool:
        """
        删除指定ID的文档
        
        Args:
            document_id: 文档ID
            
        Returns:
            是否删除成功
        """
        try:
            # 删除文档
            self.vectorstore.delete(ids=[document_id])
            

            
            return True
            
        except Exception as e:
            print(f"删除文档时出错: {str(e)}")
            return False
    
    def delete_documents_by_filter(self, filter: Dict[str, Any]) -> bool:
        """
        根据过滤条件删除文档
        
        Args:
            filter: 过滤条件
            
        Returns:
            是否删除成功
        """
        try:
            # 首先查询匹配的文档
            matching_docs = self.vectorstore.get(where=filter)
            
            if matching_docs and matching_docs['ids']:
                # 删除匹配的文档
                self.vectorstore.delete(ids=matching_docs['ids'])
                

                
            return True
            
        except Exception as e:
            print(f"根据过滤条件删除文档时出错: {str(e)}")
            return False
    
    def count_documents(self) -> int:
        """
        获取向量数据库中的文档数量
        
        Returns:
            文档数量
        """
        try:
            # 获取所有文档
            results = self.vectorstore.get()
            return len(results.get('ids', [])) if results else 0
        except Exception as e:
            print(f"获取文档数量时出错: {str(e)}")
            return 0
    
    def reset_database(self) -> bool:
        """
        重置向量数据库，删除所有文档
        
        Returns:
            是否重置成功
        """
        try:
            # 获取所有文档ID
            results = self.vectorstore.get()
            if results and results['ids']:
                # 删除所有文档
                self.vectorstore.delete(ids=results['ids'])
                

            
            return True
            
        except Exception as e:
            print(f"重置向量数据库时出错: {str(e)}")
            return False

# 提供默认实例，设置force_reset=False以保留已有数据
chromadb_manager = ChromaDBManager(force_reset=False)